from sklearn.model_selection._split import KFold
from project_settings import  network_configs,init_object,use_network_type,train_params

from tensorflow import  keras
from nlp_tools.metrics.classification import F1CategoryCallback
import numpy as np
import os
from tensorflow.keras import backend as K



def create_model():
    configs = network_configs[use_network_type]

    # 初始化分词器
    bert_tokenizer = init_object(configs, 'Tokenizer')

    # 初始化向量层embedding
    bert_embedding = init_object(configs, 'Embedding')
    bert_embedding.vocab_size=len(bert_tokenizer.token_dict)

    # 初始化句子和标签processor
    sequenceProcessor = init_object(configs, 'SentenceProcessor', object_params_init={'text_tokenizer': bert_tokenizer})
    labelProcessor = init_object(configs, 'LabelProcessor')

    # 构建模型
    model = init_object(configs, 'Network',
                        object_params_init={'embedding': bert_embedding, 'text_processor': sequenceProcessor,
                                            'label_processor': labelProcessor})
    return model


def train_and_evaluate_model(model,train_data,valid_data,model_save_path):
    early_stop = keras.callbacks.EarlyStopping(patience=20)
    ner_f1_save_callback = F1CategoryCallback(model, model_save_path, valid_data)
    model.fit(train_data, validate_data=valid_data, callbacks=[early_stop, ner_f1_save_callback],
              **train_params['fit_params'])


if __name__ == "__main__":
    n_folds = 10
    train_data = np.array(train_params['data_loader'].load_data(train_params['train_data']))
    #dev_data = train_params['data_loader'].load_data(train_params['dev_data'])
    skf = KFold( n_splits=n_folds, shuffle=True)



    for i, (train_index,test_index) in enumerate(skf.split(train_data)):
            print("Running Fold", i+1, "/", n_folds)
            model = None # Clearing the NN.
            K.clear_session()
            model = create_model()

            save_path = os.path.join(train_params['model_save_path'],"kold"+str(i))
            kold_train = train_data[train_index]
            kold_test = train_data[test_index]
            train_and_evaluate_model(model, kold_train.tolist(), kold_test.tolist(),save_path)